import yaml
import re
import json
import logging
from datetime import datetime
from mitmproxy import ctx, http
from pymongo import MongoClient
from pathlib import Path

class ArticleCaptureAddon:
    def __init__(self):
        self.config = self.load_config()
        self.setup_logging()
        self.setup_mongodb()
        # 初始化序列号跟踪字典，格式: {日期: 当前序列号}
        self.sequence_tracker = {}

    def setup_logging(self):
        # 创建日志目录
        log_dir = Path(__file__).parent / 'logs'
        log_dir.mkdir(exist_ok=True)
        
        # 日志文件名包含日期
        log_filename = log_dir / f"capture_{datetime.now().strftime('%Y%m%d')}.log"
        
        # 配置日志
        self.logger = logging.getLogger('ArticleCapture')
        self.logger.setLevel(logging.INFO)
        
        # 避免重复添加处理器
        if not self.logger.handlers:
            # 文件处理器
            file_handler = logging.FileHandler(log_filename, encoding='utf-8')
            # 控制台处理器
            console_handler = logging.StreamHandler()
            
            # 日志格式
            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            file_handler.setFormatter(formatter)
            console_handler.setFormatter(formatter)
            
            self.logger.addHandler(file_handler)
            self.logger.addHandler(console_handler)

            

    def setup_mongodb(self):
        """设置MongoDB连接"""
        try:
            db_config = self.config.get('database', {})
            host = db_config.get('host', 'localhost')
            port = db_config.get('port', 27017)
            db_name = db_config.get('db_name', 'InfoStore')
            
            self.client = MongoClient(host, port)
            self.db = self.client[db_name]
            self.logger.info(f"MongoDB连接成功: {host}:{port}/{db_name}")
        except Exception as e:
            self.logger.error(f"MongoDB连接失败: {str(e)}")
            self.client = None
            self.db = None

    def load_config(self):
        config_path = Path(__file__).parent / 'config.yaml'
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            self.logger.error(f"配置文件未找到: {config_path}")
            return {}
        except Exception as e:
            self.logger.error(f"加载配置文件失败: {str(e)}")
            return {}

    def response(self, flow: http.HTTPFlow):
        # 遍历所有规则
        for rule in self.config.get('rules', []):
            collection_name = rule.get('Collection')
            url_pattern = rule.get('url_pattern')
            
            if not collection_name or not url_pattern:
                continue
                    
            # 构建完整的URL进行匹配
            full_url = flow.request.pretty_url
            #self.logger.info(f"检查URL: {full_url}")
            #self.logger.info(f"匹配模式: {url_pattern}")
            
            # 检查URL是否匹配
            if re.search(url_pattern, full_url):
                self.logger.info(f"URL匹配成功: {full_url}")
                #self.save_response(flow, collection_name, url_pattern)
                self.process_response(flow, rule)
                break

    def process_response(self, flow: http.HTTPFlow, rule):
        """处理响应并保存到MongoDB"""
        try:
            collection_name = rule.get('Collection')
            fields = rule.get('fields', [])
            
            # 解析响应JSON
            self.logger.info(f"解析响应JSON: {collection_name}, fields: {fields}")
            try:
                response_data = flow.response.json()
            except json.JSONDecodeError:
                self.logger.error(f"无法解析JSON响应: {flow.request.path}")
                return
            
            # 提取指定字段
            extracted_data = self.extract_fields(response_data, fields)
            
            # 添加时间戳
            extracted_data['capture_time'] = datetime.now()
            extracted_data['source_url'] = flow.request.pretty_url
            
            # 保存到MongoDB
            if self.db is not None:
                collection = self.db[collection_name]
                result = collection.insert_one(extracted_data)
                self.logger.info(f"数据已保存到MongoDB: {collection_name}, ID: {result.inserted_id}")
            
            # 同时保存到文件（保持原有功能）
            self.save_response_to_file(flow, collection_name)
            
        except Exception as e:
            self.logger.error(f"处理响应失败: {str(e)}")

    def extract_fields(self, data, fields):
        """根据配置的fields提取数据"""
        extracted = {}
        
        for field in fields:
            field_name = field.get('name')
            field_path = field.get('path')
            
            if not field_name or not field_path:
                continue
                
            # 解析路径 (例如: c.dd_article_id)
            try:
                value = self.get_nested_value(data, field_path)
                if value is not None:
                    extracted[field_name] = value
            except Exception as e:
                self.logger.warning(f"提取字段失败 {field_path}: {str(e)}")
        
        return extracted

    def get_nested_value(self, data, path):
        """获取嵌套字典中的值"""
        keys = path.split('.')
        current = data
        
        for key in keys:
            if isinstance(current, dict) and key in current:
                current = current[key]
            else:
                return None
        
        return current

    def save_response_to_file(self, flow: http.HTTPFlow, collection_name):
        """保存响应到文件（保持原有功能）"""
        try:
            # 创建data目录
            data_dir = Path(__file__).parent / 'data'
            data_dir.mkdir(exist_ok=True)
            
            # 获取当前日期和序列号
            current_date = datetime.now().strftime('%Y%m%d')
            sequence = self._get_next_sequence(current_date)
            
            # 生成文件名：Collection+时间戳+两位序列号
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"{collection_name}_{timestamp}_{sequence:02d}.json"
            file_path = data_dir / filename
            
            # 保存完整响应内容
            with open(file_path, 'w', encoding='utf-8') as f:
                try:
                    # 尝试解析为JSON并格式化
                    json_data = flow.response.json()
                    json.dump(json_data, f, ensure_ascii=False, indent=2)
                except json.JSONDecodeError:
                    # 非JSON响应直接保存文本
                    f.write(flow.response.text)
                
            self.logger.info(f"已保存响应内容到: {file_path}")
            
        except Exception as e:
            self.logger.error(f"保存响应失败: {str(e)}")

    def _get_next_sequence(self, date):
        # 获取当天的下一个序列号，超过99则重置为0
        if date not in self.sequence_tracker or self.sequence_tracker[date] >= 99:
            self.sequence_tracker[date] = 0
        else:
            self.sequence_tracker[date] += 1
        return self.sequence_tracker[date]

addons = [ArticleCaptureAddon()]